"""
抛体问题求解器

date: 2025/7/16
author: SiHeng Tang
file: drop_physics.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import math

import numpy as np

G = 9.80665  # m/s²
RH_0 = 1.225  # kg/m³

# 所有的落体计算模型
ALL_MODELS = ["simple_g", "air_real"]

# 计算缓存
CALC_CACHE = {}


def _drop_backend_simple_g_model(alt: float,
                                 velocity: list):
    """
    使用最简单的初等物理计算无空气阻力的抛体运动，
    先计算落地时间，然后使用 “匀速直线运动” 推测落点。

    :param alt: 投放高度
    :param velocity: 初速度
    :return: 相对于投掷点地面投影的落点偏移
    """
    t = (- 2 * velocity[2] + math.sqrt(4 * (velocity[2] ** 2) + 8 * alt * G)) / (2 * G)
    return t * velocity[0], t * velocity[1]


def _drop_backend_air_real_model(alt: float,
                                 velocity: list,
                                 mass=0.5,
                                 cd=1.15,
                                 a=0.018,
                                 dt=0.01):
    """
    使用 RK4 积分器实现的精确落点估算函数

    :param alt: 投放高度
    :param velocity: 初始速度 (vn, ve, vd) m/s，北东地坐标系
    :param mass: 质量
    :param cd: 阻力系数
    :param a: 有效界面
    :param dt: 积分步长
    :return: 相对于投掷点地面投影的落点偏移
    """
    # 预分配轨迹数组（动态扩容）
    max_steps = int(60 / dt)  # 最长 60 s
    t = np.arange(0, max_steps * dt, dt)
    r = np.empty((max_steps, 3))
    v = np.empty_like(r)

    # 初始条件
    r[0] = [0., 0., alt]
    v[0] = np.array([velocity[0], velocity[1], -velocity[2]], dtype=float)

    def acc(rv):
        pos, vel = rv[:3], rv[3:]
        speed = np.linalg.norm(vel)
        if speed == 0:
            drag = np.zeros(3)
        else:
            Fd = 0.5 * RH_0 * cd * a * speed ** 2
            drag = -Fd * vel / speed / mass
        grav = np.array([0, 0, -G])
        return np.concatenate([vel, drag + grav])

    # RK4 循环
    for i in range(1, max_steps):
        k1 = acc(np.concatenate([r[i - 1], v[i - 1]]))
        k2 = acc(np.concatenate([r[i - 1] + 0.5 * dt * k1[:3], v[i - 1] + 0.5 * dt * k1[3:]]))
        k3 = acc(np.concatenate([r[i - 1] + 0.5 * dt * k2[:3], v[i - 1] + 0.5 * dt * k2[3:]]))
        k4 = acc(np.concatenate([r[i - 1] + dt * k3[:3], v[i - 1] + dt * k3[3:]]))

        r[i] = r[i - 1] + dt / 6 * (k1[:3] + 2 * k2[:3] + 2 * k3[:3] + k4[:3])
        v[i] = v[i - 1] + dt / 6 * (k1[3:] + 2 * k2[3:] + 2 * k3[3:] + k4[3:])

        if r[i, 2] <= 0:  # 触地
            # 线性插值求精确落点
            t_impact = t[i - 1] + dt * (0 - r[i - 1, 2]) / (r[i, 2] - r[i - 1, 2])
            x_imp = np.interp(t_impact, [t[i - 1], t[i]], [r[i - 1, 0], r[i, 0]])
            y_imp = np.interp(t_impact, [t[i - 1], t[i]], [r[i - 1, 1], r[i, 1]])
            break
    else:
        raise RuntimeError("Exceeded max simulation time")

    # 构建坐标系，注意映射关系
    return float(x_imp), float(y_imp)


def predict_drop_point(alt: float, velocity: list, model="simple_g", **kwargs):
    """
    落点预测函数，调用不同的后端使用不同的参数进行落点预测。
    落点预测暂时不是 MJ 的开发重点，目前提供两种模型可供使用::

        air_real: 使用空气阻力进行修正的落体动力模型，使用 RK4 作为微分方程数值求解器
        simple_g: 使用最简单的初等物理自由落体公式，忽略其他所有外界影响因素

    落点预测是一个核心或者将来的核心工作吗，不是。你可以继续优化，我们对这方面保持开放态度，
    但是对于过于复杂的模型考虑使用更合适的部署方式，而不是直接写在这里。

    :param alt: 高度，相对地面
    :param velocity: 速度矢量，三维空间，向下 Z 为正
    :param model: 模型选择，默认 simple_g
    :param kwargs: 模型参数，直接传给模型接口
    :return: 二维数组，代表抛体的地面投影位移矢量
    """
    if model == "simple_g":
        return _drop_backend_simple_g_model(alt, velocity)
    elif model == "air_real":
        return _drop_backend_air_real_model(alt, velocity, **kwargs)

    return None, None


if __name__ == "__main__":
    altitude = 40
    speed_vector = [10, 0, 0]

    print(f"Testing drop predict with initial parameters: {altitude=} {speed_vector=}")
    for m in ALL_MODELS:
        print(f"Model {m} predict {predict_drop_point(altitude, speed_vector, m)}")
    # 结论就是必须上制导，每个模型预测都不相同，最大差了一倍
